# %%
import serial
from serial import Serial
import struct
import numpy as np
import time



# Initialize and open serial port for data logging
try:
    ser = Serial()
    ser.baudrate = 2000000
    ser.port = 'COM18'  # Change this to your port from the list above
    ser.open()  # Add this line to open the port
    print("Serial port opened successfully")
except Exception as e:
    print(f"Error opening serial port: {e}")
    exit()  # Exit if port cannot be opened

# Initialize and open serial port for motor control
try:
    ser_motor = Serial()
    ser_motor.baudrate = 1152000
    ser_motor.port = 'COM16'  # Change this to your port from the list above
    ser_motor.open()  # Add this line to open the port
    print("Motor control serial port opened successfully")
except Exception as e:
    print(f"Error opening serial port: {e}")
    exit()  # Exit if port cannot be opened

#%% 
# start motor
ser_motor.write(b'\x03')
#%%
# setup current magnitude
ser_motor.write(b'\x06')

#%%

# Initialize data arrays
data1_array = []
data2_array = []
data3_array = []
data4_array = []
data5_array = [] 
data6_array = []  # Added 6th channel
POINTS_TO_COLLECT = 1000

try:
    while len(data1_array) < POINTS_TO_COLLECT:
        if ser.in_waiting >= 28:  # Changed to 28 bytes (7 floats * 4 bytes)
            data = ser.read(28)
            data1, data2, data3, data4, data5, data6 = struct.unpack('ffffff', data[:24])  # Updated unpacking
            data1_array.append(data1)
            data2_array.append(data2)
            data3_array.append(data3)
            data4_array.append(data4)
            data5_array.append(data5)
            data6_array.append(data6)  # Added 6th channel collection
            #print(f"Collected point {len(data1_array)}/1000")
            
except KeyboardInterrupt:
    print("Data collection interrupted")
finally:
    if ser.is_open:  # Check if port is open before closing
        ser.close()
        print("Serial port closed")

# Convert to numpy arrays
data1_array = np.array(data1_array)
data2_array = np.array(data2_array)
data3_array = np.array(data3_array)
data4_array = np.array(data4_array)
data5_array = np.array(data5_array)
data6_array = np.array(data6_array)  # Added 6th channel

print("Data collection complete")

# %%
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import numpy as np

# Create figure with shared x-axis
fig = make_subplots(
    rows=6, 
    cols=1,
    subplot_titles=('Channel 1', 'Channel 2', 'Channel 3', 'Channel 4', 'Channel 5', 'Channel 6'),
    shared_xaxes=True,  # Enable shared x-axis
    vertical_spacing=0.02  # Reduce spacing between subplots
)

# Add traces
fig.add_trace(go.Scatter(x=np.arange(len(data1_array)), y=data1_array, name="Data 1", line=dict(color='royalblue', width=2)), row=1, col=1)
fig.add_trace(go.Scatter(x=np.arange(len(data2_array)), y=data2_array, name="Data 2", line=dict(color='firebrick', width=2)), row=2, col=1)
fig.add_trace(go.Scatter(x=np.arange(len(data3_array)), y=data3_array, name="Data 3", line=dict(color='green', width=2)), row=3, col=1)
fig.add_trace(go.Scatter(x=np.arange(len(data4_array)), y=data4_array, name="Data 4", line=dict(color='orange', width=2)), row=4, col=1)
fig.add_trace(go.Scatter(x=np.arange(len(data5_array)), y=data5_array, name="Data 5", line=dict(color='brown', width=2)), row=5, col=1)
fig.add_trace(go.Scatter(x=np.arange(len(data6_array)), y=data6_array, name="Data 6", line=dict(color='purple', width=2)), row=6, col=1)

# Update layout with range slider and synchronized zoom
fig.update_layout(
    height=1080,
    showlegend=True,
    title_text="Collected Data Analysis",
    template="plotly_white",
    hovermode='x unified',
    xaxis6=dict(  # Add range slider to bottom subplot
        rangeslider=dict(visible=True),
        type="linear"
    )
)

# Update all x-axes to be linked
for i in range(1, 7):
    fig.update_xaxes(matches='x', row=i, col=1)

# Update axes labels
fig.update_xaxes(title_text="Sample Number", gridcolor='lightgrey')
fig.update_yaxes(title_text="Amplitude", gridcolor='lightgrey')

fig.show()



# %%
